from pylab import *
from scipy import *
from scipy import optimize

# if you experience problem "optimize not found", try to uncomment the following line. The problem is present at least at Ubuntu Lucid python scipy package
# from scipy import optimize

# Generate data points with noise
num_points = 500
Tx = linspace(0., 200., num_points)
Ty = Tx

tX = 1.6*cos(2*pi*(Tx/29.-1.32/360.)) +1.4*((0.5-rand(num_points))*exp(2*rand(num_points)**2))
# Fit the first set

fitfunc = lambda p, x: p[0]*cos(2*pi*(x/p[1]+p[2]/360.)) # Target function
errfunc = lambda p, x, y: fitfunc(p, x) - y # Distance to the target function
p0 = [1.6, 31., -0.] # Initial guess for the parameters
p1, success = optimize.leastsq(errfunc, p0[:], args=(Tx, tX))

time = linspace(Tx.min(), Tx.max(), 500)
plot(Tx, tX, "ro", time, fitfunc(p1, time), "r-") # Plot of the data and the fit


# Legend the plot
title("Pump-probe oscillation fitting")
xlabel("Time delay (ps)")
ylabel("$\Delta R/R$")
legend(('x position', 'x fit'))

ax = axes()

text(0.8, 0.07,
     'Periodic :  %5.3f ps.' % (p1[1]),
     fontsize=16,
     horizontalalignment='center',
     verticalalignment='center',
     transform=ax.transAxes)

show()
